import numpy as np
import scipy as sp
import scipy.integrate as integrate
import pandas as pd
import scipy.interpolate as interpol
import matplotlib.pyplot as plt
from scipy.stats.mstats import mquantiles


def cloc2DF(evalname = "evalpts.txt", outputname="output.txt"):
    #read in the evaluation positions
    df = pd.read_csv(evalname, sep=" ", skiprows = 2, header=None,skipinitialspace=True)

    ncol = len(df.columns)
    #read in the PDFs
    data = pd.read_csv(outputname,sep=" ",skiprows = 2, header=None,skipinitialspace=True)
    #combine the two
    for i in xrange(ncol):
        data['dim'+str(i)] = df[i]
    return data

if __name__ == "__main__":
    input_data = pd.read_csv("testdata0.csv", names=["x","y"])
    data = cloc2DF()
    uniq_covariate = np.unique(data.dim0)
    nuniq = len(uniq_covariate)
    qs = [0.1, 0.5, 0.9]
    quants = np.zeros([nuniq, len(qs),3])
    for ui, u in enumerate(uniq_covariate):
        sub = data[data.dim0 == u]
        norms = (np.trapz(sub, x=np.array(sub.dim1), axis=0))[0:-2]
        subb=np.array(sub)[:,0:-2]
        #cut out last 2, which are the dimensions themselves
        gridx = np.transpose(np.array(sub.dim1)*np.transpose(np.ones(subb.shape)))
        ints = integrate.cumtrapz(subb, x=gridx, axis=0)
        x_axis = sub.dim1[1:] # integration gives 1 less data point
	# so I interpret it as being shifted by one
        normed = ints/norms
        nreals = normed.shape[1]
        q = np.zeros([nreals, len(qs)])
        for i in xrange(nreals):
            f = interpol.interp1d(normed[:,i], x_axis)
            for iqi, qi in enumerate(qs):
                q[i,iqi] = f(qi)
        #quants[ui,:] = np.median(q,0)
        quants[ui,:] = np.transpose(mquantiles(q, prob=qs, axis=0))
        #to get the first pdf, use normed[:,0].shape
    #df = pd.DataFrame(

    #median plot
    fig = plt.figure(figsize=(20,6))
    ax1 = fig.add_subplot(1,3,2)
    #ax1.plot(input_data.x, input_data.y, 'o', color="#AAAAAA", alpha=0.1, mec='none')
    ax1.plot(uniq_covariate, uniq_covariate, color='red')
    #median of the median
    ax1.plot(uniq_covariate, quants[:,1,1], color='black')
    #0.1 quantiles of the median
    ax1.plot(uniq_covariate, quants[:,1,0], color='black', ls='dotted')
    #0.9 quantiles of the median
    ax1.plot(uniq_covariate, quants[:,1,2], color='black', ls='dotted')
    ax1.scatter(input_data.x, input_data.y, marker='o', color="#AAAAAA", alpha=0.1, edgecolors="none")
    ax1.set_title("Median")
    ax1.set_ylim([-5,5])
    ax1.set_xlim([-3,3])
    #------------------
    ax2 = fig.add_subplot(1,3,1)
    ax2.scatter(input_data.x, input_data.y, marker='o', color="#AAAAAA", alpha=0.1, edgecolors="none")
    ax2.plot(uniq_covariate, uniq_covariate-1.281552, color='red')
    #median of the 0.1 quantile
    ax2.plot(uniq_covariate, quants[:,0,1], color='black')
    #0.1 quantiles of the 0.1 quantile
    ax2.plot(uniq_covariate, quants[:,0,0], color='black', ls='dotted')
    #0.9 quantiles of the 0.1 quantile
    ax2.plot(uniq_covariate, quants[:,0,2], color='black', ls='dotted')
    ax2.set_title("0.1 Quantile")
    ax2.set_ylim([-5,5])
    ax2.set_xlim([-3,3])
    #------------------
    ax3 = fig.add_subplot(1,3,3)
    ax3.scatter(input_data.x, input_data.y, marker='o', color="#AAAAAA", alpha=0.1, edgecolors="none")
    ax3.plot(uniq_covariate, uniq_covariate+1.281552, color='red')
    #median of the 0.9 quantile
    ax3.plot(uniq_covariate, quants[:,2,1], color='black')
    #0.1 quantiles of the 0.9 quantile
    ax3.plot(uniq_covariate, quants[:,2,0], color='black', ls='dotted')
    #0.9 quantiles of the 0.9 quantile
    ax3.plot(uniq_covariate, quants[:,2,2], color='black', ls='dotted')
    ax3.set_title("0.9 Quantile")
    ax3.set_ylim([-5,5])
    ax3.set_xlim([-3,3])
    plt.show()
    
    
